#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Agent自动更新模块
通过定期检查控制端的版本信息，自动下载和更新Agent代码
"""

import os
import sys
import json
import time
import hashlib
import shutil
import subprocess
import requests
from pathlib import Path
from datetime import datetime
from typing import Optional, Dict, Any
from logger_setup import logger


class AutoUpdater:
    """Agent自动更新器"""
    
    def __init__(self, config_file: str = "config.json"):
        """
        初始化自动更新器
        
        Args:
            config_file: 配置文件路径
        """
        self.config_file = config_file
        self.config = self._load_config()
        
        # 从配置中获取更新相关参数
        self.update_url = self.config.get('update_url', '')
        self.check_interval = self.config.get('update_check_interval', 3600)  # 默认1小时检查一次
        self.auto_update_enabled = self.config.get('auto_update_enabled', True)
        
        # 版本文件路径
        self.version_file = Path(__file__).parent.parent / 'version.json'
        self.current_version = self._load_current_version()
        
        # 下载目录
        self.download_dir = Path(__file__).parent.parent / 'downloads'
        self.download_dir.mkdir(exist_ok=True)
        
        # 备份目录
        self.backup_dir = Path(__file__).parent.parent / 'backups'
        self.backup_dir.mkdir(exist_ok=True)
        
        logger.info(f"自动更新器初始化完成，当前版本: {self.current_version}")
    
    def _load_config(self) -> Dict[str, Any]:
        """加载配置文件"""
        try:
            config_path = Path(__file__).parent.parent / self.config_file
            if config_path.exists():
                with open(config_path, 'r', encoding='utf-8') as f:
                    return json.load(f)
        except Exception as e:
            logger.error(f"加载配置文件失败: {e}")
        
        return {}
    
    def _load_current_version(self) -> str:
        """加载当前版本信息"""
        try:
            if self.version_file.exists():
                with open(self.version_file, 'r', encoding='utf-8') as f:
                    version_info = json.load(f)
                    return version_info.get('version', '1.0.0')
        except Exception as e:
            logger.error(f"加载版本信息失败: {e}")
        
        return '1.0.0'
    
    def _save_version(self, version: str, build_time: str = None):
        """保存版本信息"""
        try:
            version_info = {
                'version': version,
                'build_time': build_time or datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                'update_time': datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            }
            
            with open(self.version_file, 'w', encoding='utf-8') as f:
                json.dump(version_info, f, indent=2, ensure_ascii=False)
            
            logger.info(f"版本信息已保存: {version}")
            return True
        except Exception as e:
            logger.error(f"保存版本信息失败: {e}")
            return False
    
    def check_for_updates(self) -> Optional[Dict[str, Any]]:
        """
        检查是否有新版本
        
        Returns:
            如果有新版本，返回版本信息字典，否则返回None
        """
        if not self.update_url:
            logger.warning("未配置update_url，无法检查更新")
            return None
        
        try:
            logger.info("检查更新...")
            response = requests.get(
                f"{self.update_url}/version",
                timeout=10
            )
            
            if response.status_code == 200:
                version_info = response.json()
                latest_version = version_info.get('version', '1.0.0')
                
                logger.info(f"当前版本: {self.current_version}, 最新版本: {latest_version}")
                
                # 比较版本号
                if self._compare_versions(latest_version, self.current_version) > 0:
                    logger.info(f"发现新版本: {latest_version}")
                    return version_info
                else:
                    logger.info("当前已是最新版本")
                    return None
            else:
                logger.error(f"检查更新失败，HTTP状态码: {response.status_code}")
                return None
                
        except Exception as e:
            logger.error(f"检查更新异常: {e}")
            return None
    
    def _compare_versions(self, version1: str, version2: str) -> int:
        """
        比较两个版本号
        
        Args:
            version1: 版本号1
            version2: 版本号2
        
        Returns:
            1: version1 > version2
            0: version1 == version2
            -1: version1 < version2
        """
        try:
            v1_parts = [int(x) for x in version1.split('.')]
            v2_parts = [int(x) for x in version2.split('.')]
            
            # 补齐长度
            max_len = max(len(v1_parts), len(v2_parts))
            v1_parts.extend([0] * (max_len - len(v1_parts)))
            v2_parts.extend([0] * (max_len - len(v2_parts)))
            
            for i in range(max_len):
                if v1_parts[i] > v2_parts[i]:
                    return 1
                elif v1_parts[i] < v2_parts[i]:
                    return -1
            
            return 0
        except Exception as e:
            logger.error(f"版本比较失败: {e}")
            return 0
    
    def download_update(self, version_info: Dict[str, Any]) -> Optional[str]:
        """
        下载更新包
        
        Args:
            version_info: 版本信息
        
        Returns:
            下载成功返回文件路径，失败返回None
        """
        try:
            version = version_info.get('version')
            download_url = version_info.get('download_url')
            file_hash = version_info.get('md5')
            
            if not download_url:
                logger.error("下载URL为空")
                return None
            
            # 下载文件
            logger.info(f"开始下载更新包: {download_url}")
            response = requests.get(download_url, stream=True, timeout=300)
            
            if response.status_code == 200:
                file_path = self.download_dir / f"agent_{version}.tar.gz"
                
                with open(file_path, 'wb') as f:
                    for chunk in response.iter_content(chunk_size=8192):
                        if chunk:
                            f.write(chunk)
                
                logger.info(f"下载完成: {file_path}")
                
                # 验证文件哈希
                if file_hash and not self._verify_file_hash(file_path, file_hash):
                    logger.error("文件校验失败")
                    file_path.unlink()
                    return None
                
                return str(file_path)
            else:
                logger.error(f"下载失败，HTTP状态码: {response.status_code}")
                return None
                
        except Exception as e:
            logger.error(f"下载更新包异常: {e}")
            return None
    
    def _verify_file_hash(self, file_path: Path, expected_hash: str) -> bool:
        """验证文件哈希值"""
        try:
            md5 = hashlib.md5()
            with open(file_path, 'rb') as f:
                for chunk in iter(lambda: f.read(4096), b''):
                    md5.update(chunk)
            
            actual_hash = md5.hexdigest()
            
            if actual_hash.lower() == expected_hash.lower():
                logger.info("文件校验成功")
                return True
            else:
                logger.error(f"文件校验失败，期望: {expected_hash}, 实际: {actual_hash}")
                return False
        except Exception as e:
            logger.error(f"文件校验异常: {e}")
            return False
    
    def backup_current_version(self) -> bool:
        """备份当前版本"""
        try:
            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
            backup_path = self.backup_dir / f"agent_{self.current_version}_{timestamp}"
            
            logger.info(f"备份当前版本到: {backup_path}")
            
            # 获取当前agent目录
            agent_dir = Path(__file__).parent.parent
            
            # 排除不需要备份的目录和文件
            exclude_dirs = {'downloads', 'backups', 'logs', '__pycache__', '.git'}
            
            # 创建备份
            shutil.copytree(
                agent_dir,
                backup_path,
                ignore=lambda dir, files: [
                    f for f in files 
                    if f in exclude_dirs or f.endswith('.pyc')
                ]
            )
            
            logger.info("备份成功")
            return True
            
        except Exception as e:
            logger.error(f"备份失败: {e}")
            return False
    
    def apply_update(self, update_file: str, version_info: Dict[str, Any]) -> bool:
        """
        应用更新
        
        Args:
            update_file: 更新包文件路径
            version_info: 版本信息
        
        Returns:
            更新成功返回True，否则返回False
        """
        try:
            logger.info(f"开始应用更新: {update_file}")
            
            # 1. 备份当前版本
            if not self.backup_current_version():
                logger.error("备份失败，取消更新")
                return False
            
            # 2. 解压更新包
            agent_dir = Path(__file__).parent.parent
            
            logger.info("解压更新包...")
            import tarfile
            with tarfile.open(update_file, 'r:gz') as tar:
                tar.extractall(agent_dir)
            
            # 3. 更新版本信息
            new_version = version_info.get('version')
            build_time = version_info.get('build_time')
            self._save_version(new_version, build_time)
            
            # 4. 执行更新后脚本（如果存在）
            update_script = agent_dir / 'update_post.sh'
            if update_script.exists():
                logger.info("执行更新后脚本...")
                subprocess.run(['bash', str(update_script)], check=True)
            
            logger.info(f"更新成功，版本: {new_version}")
            
            # 5. 删除下载的更新包
            Path(update_file).unlink()
            
            return True
            
        except Exception as e:
            logger.error(f"应用更新失败: {e}")
            return False
    
    def restart_agent(self):
        """重启Agent进程"""
        try:
            logger.info("准备重启Agent...")
            
            # 获取当前进程的启动脚本
            agent_dir = Path(__file__).parent.parent
            start_script = agent_dir / 'start.sh'
            
            if not start_script.exists():
                logger.error("启动脚本不存在")
                return False
            
            # 使用nohup在后台重启
            logger.info("重启Agent进程...")
            subprocess.Popen(
                ['nohup', 'bash', str(start_script)],
                cwd=str(agent_dir),
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL
            )
            
            # 给新进程一点启动时间
            time.sleep(2)
            
            # 退出当前进程
            logger.info("退出当前进程...")
            sys.exit(0)
            
        except Exception as e:
            logger.error(f"重启Agent失败: {e}")
            return False
    
    def run_update_check(self):
        """运行更新检查（主循环）"""
        if not self.auto_update_enabled:
            logger.info("自动更新已禁用")
            return
        
        logger.info("自动更新器启动")
        
        while True:
            try:
                # 检查更新
                version_info = self.check_for_updates()
                
                if version_info:
                    # 发现新版本，下载更新
                    update_file = self.download_update(version_info)
                    
                    if update_file:
                        # 应用更新
                        if self.apply_update(update_file, version_info):
                            logger.info("更新应用成功，准备重启...")
                            
                            # 等待5秒后重启
                            time.sleep(5)
                            self.restart_agent()
                        else:
                            logger.error("应用更新失败")
                
                # 等待下次检查
                logger.info(f"等待 {self.check_interval} 秒后再次检查更新...")
                time.sleep(self.check_interval)
                
            except Exception as e:
                logger.error(f"更新检查循环异常: {e}")
                time.sleep(60)  # 出错后等待1分钟再试


def main():
    """测试入口"""
    updater = AutoUpdater()
    updater.run_update_check()


if __name__ == '__main__':
    main()

